import os
import torch
import random
import logging
import argparse
import numpy as np
import utils_data as utils
import warnings
from fastapi import FastAPI
import torch
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from fastapi import Request
import uvicorn
import sys
sys.path.append(os.path.dirname(sys.path[0]))
logging.basicConfig(level=logging.DEBUG)

from metrics import get_entities
from data_loader import DataLoader
from SequenceTagger import BertForSequenceTagging
from detect_text_error import process_medical_report, detected_conflict

app = FastAPI()
app.mount("/static", StaticFiles(directory="./static"), name="static")

@app.get("/")
async def read_index():
    return FileResponse(os.path.join('./templates', 'index.html'))

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', default='radio', help="Directory containing the dataset")
parser.add_argument('--seed', type=int, default=2025, help="random seed for initialization")
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
warnings.filterwarnings("ignore", category=UserWarning)

def interAct(model, data_iterator, params, mark='Interactive', verbose=False):
    """Evaluate the model on `steps` batches."""
    # set model to evaluation mode
    model.eval()
    idx2tag = params.idx2tag

    batch_data, batch_token_starts = next(data_iterator)
    batch_masks = batch_data.gt(0)
        
    batch_output = model((batch_data, batch_token_starts), token_type_ids=None, attention_mask=batch_masks)[0]  # shape: (batch_size, max_len, num_labels)
    batch_output = batch_output.detach().cpu().numpy()
    
    pred_tags = []
    pred_tags.extend([[idx2tag.get(idx) for idx in indices] for indices in np.argmax(batch_output, axis=2)])
    
    return(get_entities(pred_tags))

def bert_ner_init():
    """"模型和数据加载器初始化"""
    args = parser.parse_args()
    tagger_model_dir = './experiments/' + args.dataset

    json_path = os.path.join(tagger_model_dir, 'params.json')
    assert os.path.isfile(json_path), "No json configuration file found at {}".format(json_path)
    params = utils.Params(json_path)

    params.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    random.seed(args.seed)
    torch.manual_seed(args.seed)
    params.seed = args.seed

    utils.set_logger(os.path.join(tagger_model_dir, 'evaluate.log'))
    logging.info("Loading the dataset...")
    data_dir = './data/' + args.dataset
    bert_class = './experiments/bert-base-chinese'
    data_loader = DataLoader(data_dir, bert_class, params, token_pad_idx=0, tag_pad_idx=-1)
    try :
        model = BertForSequenceTagging.from_pretrained(tagger_model_dir)
    except Exception as e:
        logging.error(f"加载模型失败: {str(e)}")
        raise
    model.to(params.device)
    return model, data_loader, args.dataset, params

def BertNerResponse(model, queryString):
    model, data_loader, dataset, params = model
    queryString = [i for i in queryString]  # bio格式时，使用这个语句
    # queryString = nltk.word_tokenize(queryString)  # bmes格式时使用这个
    with open('data/' + dataset + '/interactive/sentences.txt', 'w', encoding='utf-8') as f:
        f.write(' '.join(queryString))

    inter_data = data_loader.load_data('interactive')
    inter_data_iterator = data_loader.data_iterator(inter_data, shuffle=False)
    result = interAct(model, inter_data_iterator, params)
    res = []
    for item in result:
        res.append((' '.join(queryString[item[1]:item[2] + 1]), item[0]))
    return res

def get_bert_ner_service():
    """获取BERT-NER服务实例（单例模式）"""
    if not hasattr(app, 'model'):
        app.model, app.data_loader, app.dataset, app.params = bert_ner_init()
    return app.model, app.data_loader, app.dataset, app.params

def process_text_with_bert(text):
    """
    封装BERT-NER处理文本的API方法
    :param text: 要处理的文本
    :return: 处理结果列表，格式为[(实体文本, 实体类型), ...]
    """
    try:
        model, data_loader, dataset, params = get_bert_ner_service()
        sequence = BertNerResponse((model, data_loader, dataset, params), text)
        return [(t[0].replace(' ', ''), t[1].replace(' ', '')) for t in sequence]
    except Exception as e:
        logging.error(f"BERT-NER处理失败: {str(e)}")
        return []

@app.post("/sentence_process")
async def sentence_process(request: Request):  # 改为异步并接收Request对象
    """处理API请求的主入口"""
    try:
        # 添加内容类型验证
        content_type = request.headers.get('Content-Type')
        if content_type != 'application/json':
            return {"status": "error", "message": "Invalid content type"}

        get_bert_ner_service()
        data = await request.json() # 获取请求参数
        query_imaging = data.get('imaging', '').replace('\n', ' ').replace('\r', '').strip()
        query_diagnosis = data.get('diagnosis', '').replace('\n', ' ').replace('\r', '').strip()
        imaging_data = process_text_with_bert(query_imaging)
        diagnosis_data = process_text_with_bert(query_diagnosis)

        # 分别处理两类冲突，一类是位置，一类是语义冲突
        result_organ = detected_conflict(imaging_data, diagnosis_data)
        print("-" * 110)
        for idx, item in enumerate(result_organ, 1):
            print(item)
        report_result = process_medical_report(imaging_data, diagnosis_data)
        print("-" * 110)
        for idx, item in enumerate(report_result, 1):
            print(
                f"{str(idx):<4} | {str(item.get('解剖部位', '')):<14} |"
                f"{str(item.get('影像异常', '')):<18} |"
                f"{str(item.get('诊断异常', '')):<18} |"
                f"{str(item.get('状态', '')):<18} |"
                f"{str(item.get('详情', ''))}"
            )
        
        # 确保返回的数据是JSON可序列化的
        return {
            "status": "success",
            "imaging_entities": imaging_data if imaging_data is not None else [],
            "diagnosis_entities": diagnosis_data if diagnosis_data is not None else [],
            "conflicts": [str(item) for item in result_organ] if result_organ else [],
            "report": [
                {
                    "解剖部位": str(item.get('解剖部位', '')),
                    "影像异常": str(item.get('影像异常', '')),
                    "诊断异常": str(item.get('诊断异常', '')),
                    "状态": str(item.get('状态', '')),
                    "详情": str(item.get('详情', ''))
                } 
                for item in (report_result if report_result else [])
            ]
        }
    except Exception as e:
        return {"status": "error", "message": str(e)}
    
def test_api():
    """API测试方法"""
    from fastapi.testclient import TestClient
    client = TestClient(app)
    
    # 测试正常请求
    test_data = {
        "imaging": "右肺中叶见斑片状高密度影，边界模糊",
        "diagnosis": "左肺中叶炎症"
    }
    response = client.post("/sentence_process", json=test_data)
    print("\n测试结果：")
    print(f"状态码: {response.status_code}")
    print(f"响应内容: {response.json()}")

if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8300)
    # test_api()
